#include <pthread.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>

#define __USE_GNU
#include <fcntl.h>
#undef __USE_GNU

#include "config.h"

#include "queue.h"

#include "psd_tpool.h"
#include "psd_logger.h"
#include "psd_xfer.h"
#include "psd_conn_tcp.h"
#include "psd_modules.h"
#include "psd_backend_handler.h"

#include "compat.h"
#include "splice.h"

#include <event.h>

#define SPLICE_SIZE (64*1024)
#define PSD_BUF_SIZE 1048576

#define RECV_SIDE_DONE     0x01
#define SEND_SIDE_DONE     0x02

typedef struct libeventCtrl {
    int recv_fd;
    int recvptr;
    int send_fd;
    int sendptr;

    int flags;

    struct event *r_ev;
    struct event *s_ev;

    psdXferArgs *args;

    pthread_mutex_t lock;
    pthread_cond_t cond;
} psdEvCtl;

int ev_ctl_pipe[2];
struct event ev_ctl_event;

int use_splice = 0;

static TAILQ_HEAD(listhead, event) ev_updates;
pthread_mutex_t ev_updates_lock;

int psd_xfer_libevent_init();
int __psd_xfer_libevent_init();
void psd_event_add(struct event *ev);
void *psd_event_run(void *arg);
void psd_update_events(int fd, short event, void *arg);
psdEvCtl *psd_alloc_psdevctl();
void psd_free_evctl(psdEvCtl *ctl);
void psd_new_event_handler(int fd, short event, void *arg);
psdEvCtl *psd_alloc_evctl_xfer_args(psdXferArgs *args);

int psd_xfer_libevent_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt);
int psd_xfer_libevent_splice_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt);

psdModule psd_xfer_libevent_module = {
	.desc = "libevent Module",
	.dependencies = "tcp",
        .init = psd_xfer_libevent_init
};

psdBackendHandler psd_xfer_libevent_backend_handler = {
        .name = "libevent",
        .run = psd_xfer_libevent_main
};

psdBackendHandler psd_xfer_libevent_splice_backend_handler = {
        .name = "libevent-splice",
        .run = psd_xfer_libevent_splice_main
};

psdModule *module_info() {
        return &psd_xfer_libevent_module;
}

int psd_xfer_libevent_init() {
	if (__psd_xfer_libevent_init() != 0)
		return -1;
	
	if (psd_add_backend_handler(&psd_xfer_libevent_backend_handler) != 0)
		return -1;

        return psd_add_backend_handler(&psd_xfer_libevent_splice_backend_handler);
}

int __psd_xfer_libevent_init() {
    event_init();

    if (pthread_mutex_init(&ev_updates_lock, 0))
        goto error_exit;

    if (pipe(ev_ctl_pipe) < 0)
        goto error_exit2;

    event_set(&ev_ctl_event, ev_ctl_pipe[0], EV_READ | EV_PERSIST, psd_update_events, NULL);

    TAILQ_INIT(&ev_updates);

    psd_tpool_exec(psd_event_run, NULL);

    return 0;

error_exit2:
    pthread_mutex_destroy(&ev_updates_lock);
error_exit:
    return -1;
}

int psd_xfer_libevent_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt) {
	return psd_libevent_main_loop(p_args, c_args, 0);
}

int psd_xfer_libevent_splice_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt) {
	return psd_libevent_main_loop(p_args, c_args, 1);
}

int psd_libevent_main_loop(psdXferArgs *p_args, psdXferArgs *c_args, int opt) {
	psdEvCtl *ev_ctl0, *ev_ctl1;

	if (opt) {
		use_splice = 1;

		p_args->bufsize = SPLICE_SIZE;
		c_args->bufsize = SPLICE_SIZE;
	}
	else {
		p_args->buf = malloc(sizeof(char) * p_args->bufsize);
		if (!p_args->buf)
			goto error_exit;
		
		c_args->buf = malloc(sizeof(char) * c_args->bufsize);
		if (!c_args->buf)
			goto error_exit;
	}
	
	ev_ctl0 = psd_alloc_evctl_xfer_args(p_args);
	if (!ev_ctl0) {
		psd_err(5, "couldn't allocate transfer structure");
		goto error_exit;
	}
	
	ev_ctl1 = psd_alloc_evctl_xfer_args(c_args);
	if (!ev_ctl1) {
		psd_err(5, "couldn't allocate transfer structure");
		goto error_exit_ev_ctl0;
	}
	
	psd_info(0, "Starting up %d->%d", ev_ctl0->recv_fd, ev_ctl0->send_fd);
	psd_event_add(ev_ctl0->r_ev);
	psd_info(0, "Starting up %d->%d", ev_ctl1->recv_fd, ev_ctl1->send_fd);
	psd_event_add(ev_ctl1->r_ev);

	pthread_mutex_lock(&(ev_ctl0->lock));
	if (!(ev_ctl0->flags & SEND_SIDE_DONE)) {
		pthread_cond_wait(&(ev_ctl0->cond), &(ev_ctl0->lock)); 
	}
	pthread_mutex_unlock(&(ev_ctl0->lock));
	
	pthread_mutex_lock(&(ev_ctl1->lock));
	if (!(ev_ctl1->flags & SEND_SIDE_DONE)) {
		pthread_cond_wait(&(ev_ctl1->cond), &(ev_ctl1->lock)); 
	}
	pthread_mutex_unlock(&(ev_ctl1->lock));

	psd_free_evctl(ev_ctl0);
	psd_free_evctl(ev_ctl1);

	return 0;

error_exit_ev_ctl0:
	psd_free_evctl(ev_ctl0);
error_exit:
	return -1;
}

psdEvCtl *psd_alloc_evctl_xfer_args(psdXferArgs *args) {
    psdConn *curr_src = args->src;
    psdConn *curr_sink = args->sink;
    psdConn_tcpData *src_tcp_info = curr_src->conn_private;
    psdConn_tcpData *sink_tcp_info = curr_sink->conn_private;
    psdEvCtl *ev_ctl;

    if (strcmp(curr_src->protocol, "TCP") != 0) {
        psd_err(0, "specified sink is not TCP");
        return NULL;
    }

    if (strcmp(curr_sink->protocol, "TCP") != 0) {
        psd_err(0, "specified sink is not TCP");
        return NULL;
    }

    ev_ctl = psd_alloc_psdevctl();
    if (!ev_ctl) {
        psd_err(0, "couldn't allocate evctl");
        return NULL;
    }

    ev_ctl->flags = 0;
    ev_ctl->recvptr = 0;
    ev_ctl->sendptr = 0;
    ev_ctl->recv_fd = src_tcp_info->sd;
    ev_ctl->send_fd = sink_tcp_info->sd;
    ev_ctl->args = args;

    // setup pipe for splice
    if (use_splice) {
	    int ret = pipe(args->pipes);
	    
	    if (ret < 0) {
		    psd_err(0, "couldn't create pipe: %s\n", strerror(errno));
		    return NULL;
	    }
    }

    // setup the events
    event_set(ev_ctl->r_ev, ev_ctl->recv_fd, EV_READ | EV_PERSIST, psd_new_event_handler, ev_ctl);
    event_set(ev_ctl->s_ev, ev_ctl->send_fd, EV_WRITE | EV_PERSIST, psd_new_event_handler, ev_ctl);

    return ev_ctl;
}

void psd_event_add(struct event *ev) {
    char junk;

//    psd_info(0, "adding event");

    pthread_mutex_lock(&ev_updates_lock);
    {
        TAILQ_INSERT_TAIL(&ev_updates, ev, ev_next);

        write(ev_ctl_pipe[1], &junk, 1);
    }
    pthread_mutex_unlock(&ev_updates_lock);
}

void *psd_event_run(void *arg) {
    event_add(&ev_ctl_event, NULL);

//    psd_info(0, "spawning dispatch thread");

    event_dispatch();
}

void psd_update_events(int fd, short event, void *arg) {
    struct event *ev;
    char buf[5];
    int count;

//    psd_info(0, "reading in new event");

    pthread_mutex_lock(&ev_updates_lock);
    {
        count = 0;

        while((ev = ev_updates.tqh_first) != NULL) {
            TAILQ_REMOVE(&ev_updates, ev, ev_next);

//            psd_info(0, "Adding event to queue");

            event_add(ev, NULL);
            count++;
        }

        while(count > 0) {
            int amt_to_read;

            amt_to_read = (count > sizeof(buf))?sizeof(buf):count;
            read(ev_ctl_pipe[0], buf, amt_to_read);
            count -= amt_to_read;

        }
    }
    pthread_mutex_unlock(&ev_updates_lock);
}

////////////////////////////////////////////////
// Bidirectional Support for Events
////////////////////////////////////////////////

psdEvCtl *psd_alloc_psdevctl() {
    psdEvCtl *ret_ctl;
    struct event *send_ev, *recv_ev;

    // allocate an event controller
    ret_ctl = (psdEvCtl *) malloc(sizeof(psdEvCtl));
    if (!ret_ctl)
        goto error_exit;

    // allocate a send event
    send_ev = (struct event *) malloc(sizeof(struct event));
    if (!send_ev)
        goto error_exit1;

    // allocate a recv event
    recv_ev = (struct event *) malloc(sizeof(struct event));
    if (!recv_ev)
        goto error_exit2;

    pthread_mutex_init(&ret_ctl->lock, NULL);
    pthread_cond_init(&ret_ctl->cond, NULL);

    // assign them to the controller
    ret_ctl->r_ev = recv_ev;
    ret_ctl->s_ev = send_ev;

    return ret_ctl;

error_exit2:
    free(send_ev);
error_exit1:
    free(ret_ctl);
error_exit:
    return NULL;
}

void psd_free_evctl(psdEvCtl *ctl) {
    free(ctl->r_ev);
    free(ctl->s_ev);
    free(ctl);
}

void psd_new_event_handler(int fd, short event, void *arg) {
	psdEvCtl *ev_ctl = arg;
	psdXferArgs *args = ev_ctl->args;

	int recvd;
	int sent;
	
	//    psd_info(0, "Got event: %d %d: %d->%d", fd, event, ev_ctl->recv_fd, ev_ctl->send_fd);
	
	if (event & EV_READ) {
		//psd_info(0, "Got read event: %d %d: %d->%d", fd, event, ev_ctl->recv_fd, ev_ctl->send_fd);
		
		// figure out how much space we have to receive into
		int diff = args->bufsize - ev_ctl->recvptr;
		
		//psd_info(0, "Diff : %d %d %d: %d->%d", diff, args->bufsize, ev_ctl->recvptr, ev_ctl->recv_fd, ev_ctl->send_fd);
		
		if (diff > 0) {
			// if we have space, try and read into it
			
			if (use_splice)
				recvd = ssplice(fd, NULL, args->pipes[1], NULL, diff, SPLICE_F_NONBLOCK);
			else
				recvd = recv(fd, args->buf + ev_ctl->recvptr, diff, MSG_DONTWAIT);
		       
			//psd_info(11, "Recv returned %d: %d->%d", recvd, ev_ctl->recv_fd, ev_ctl->send_fd);
			
			if (recvd > 0) {
				int prev_diff = ev_ctl->recvptr - ev_ctl->sendptr;
				ev_ctl->recvptr += recvd;
				
				// if we had no data before, we didn't want any "writable" events before
				if (prev_diff == 0) {
					event_add(ev_ctl->s_ev, NULL);
				}
				
				// if we have no space left in the buffer, delete the
				// read listener
				if (ev_ctl->recvptr >= args->bufsize) {
					event_del(ev_ctl->r_ev);
				}
			} else if (errno == EAGAIN) {
				// do nothing
				errno = 0;
			} else {
				// if the other side closed, we close too
				if (recvd < 0)
					psd_info(11, "psd_event_handler(): read(fd, pipe->buf, diff) failed: %s", strerror(errno));
				else
					psd_info(11, "psd_event_handler(): closing and exiting for %d", fd);
				
				event_del(ev_ctl->r_ev);
				
				// shutdown the read side since we may still have
				// a writer
				if (use_splice)
					close(args->pipes[1]);

				psd_conn_shutdown(ev_ctl->args->src, PSD_RECV_SIDE);
				
				// set the flag so we know to close the connection
				ev_ctl->flags |= RECV_SIDE_DONE;
				
				// wake up the sender side in case we weren't paying attention to his events
				event_add(ev_ctl->s_ev, NULL);
			}
		}
	} else if (event & EV_WRITE) {
		int diff, new_diff;
		
		diff = ev_ctl->recvptr - ev_ctl->sendptr;
		
		//psd_info(0, "Got write event: %d %d: %d->%d", fd, diff, ev_ctl->recv_fd, ev_ctl->send_fd);
		
		if (diff > 0) {

			if (use_splice)
				sent = ssplice(args->pipes[0], NULL, fd, NULL, diff, SPLICE_F_NONBLOCK);
			else
				sent = send(fd, args->buf + ev_ctl->sendptr, diff, MSG_DONTWAIT);

			//psd_info(11, "Send returned %d: %d->%d", sent, ev_ctl->recv_fd, ev_ctl->send_fd);
			
			// if len == 0?
			
			if (sent > 0) {
				ev_ctl->sendptr += sent;
			} else if (errno == EAGAIN) {
				// do nothing
				errno = 0;
			} else {
				if (sent == 0) {
					psd_info(0, "Other side closed, shutdown down: %d->%d", ev_ctl->recv_fd, ev_ctl->send_fd);
				} else {
					psd_info(0, "Received an error, shutdown down: %d->%d: %s", ev_ctl->recv_fd, ev_ctl->send_fd, strerror(errno));
				}
				
				event_del(ev_ctl->r_ev);
				event_del(ev_ctl->s_ev);

				if (use_splice) {
					close(args->pipes[0]);
					close(args->pipes[1]);
				}

				psd_conn_shutdown(ev_ctl->args->sink, PSD_SEND_SIDE);
				psd_conn_shutdown(ev_ctl->args->src, PSD_RECV_SIDE);
				
				pthread_mutex_lock(&(ev_ctl->lock));
				{
					ev_ctl->flags |= SEND_SIDE_DONE;
				}
				pthread_mutex_unlock(&(ev_ctl->lock));
				
				pthread_cond_signal(&(ev_ctl->cond));
			}
		}
		
		// if we've caught up to the receiver, go back to the beginning
		if (ev_ctl->sendptr == ev_ctl->recvptr) {
			// if the reciever has been waiting for us to catch up, "wake him up"
			if (ev_ctl->recvptr >= args->bufsize) {
				//                psd_info(0, "Waking up receiver: %d->%d", ev_ctl->recv_fd, ev_ctl->send_fd);
				event_add(ev_ctl->r_ev, NULL);
			}
			
			ev_ctl->sendptr = 0;
			ev_ctl->recvptr = 0;
		}
		
		// calculate the size of the data left to be sent
		new_diff = ev_ctl->recvptr - ev_ctl->sendptr;
		
		// if we have nothing left to send and the receiver is done, close the sender side
		if (new_diff == 0 && ev_ctl->flags & RECV_SIDE_DONE) {
			psd_info(0, "Closing the send side: %d->%d", ev_ctl->recv_fd, ev_ctl->send_fd);
			
			event_del(ev_ctl->s_ev);
			
			if (use_splice)
				close(args->pipes[0]);
			
			psd_conn_shutdown(ev_ctl->args->sink, PSD_SEND_SIDE);
			
			pthread_mutex_lock(&(ev_ctl->lock));
			{
				ev_ctl->flags |= SEND_SIDE_DONE;
			}
			pthread_mutex_unlock(&(ev_ctl->lock));
			
			pthread_cond_signal(&(ev_ctl->cond));
		} else if (new_diff == 0) {
			// if we have nothing left to send, delete the sending
			// event
			event_del(ev_ctl->s_ev);
		}
	}
}

